from typing import List, Any, Tuple

from chromadb import GetResult
from langchain_core.documents import Document


# 序列化查询结果
# @param get_result 向量数据库中查询的结果
# @return 序列化后的结果
def _get_result_to_documents(get_result: GetResult) -> List[Document]:
    if get_result is None or not isinstance(get_result) or (not get_result['documents']):
        return []
    _metadatas = get_result['metadatas'] if get_result['metadatas'] else [{}] * len(get_result['documents'])
    document_list = []
    for page_content, metadata in zip(get_result['documents'], _metadatas):
        document_list.append(Document(**{'page_content': page_content, 'metadata': metadata}))

    return document_list


# 序列化匹配搜索
# @param results 匹配搜索结果
# @return 序列化的结果
def _results_to_docs_and_scores(results: Any) -> List[Tuple[Document, float]]:
    """
    from langchain_community.vectorstores.chroma import Chroma
    """
    return [
        # TODO: Chroma can do batch querying,
        (Document(page_content=result[0], metadata=result[1] or {}), result[2])
        for result in zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        )
    ]
